import rclpy
import cv2
import cv_bridge
import numpy as np
from rclpy.node import Node
from sensor_msgs.msg import CompressedImage
from origincar_msg.msg import Sign
from geometry_msgs.msg import Twist
from std_msgs.msg import Int32
from ai_msgs.msg import PerceptionTargets
from time import time

class Follower(Node):
    def __init__(self):
        super().__init__('line_follower')
        self.get_logger().info("Start line follower.")
        
        self.bridge = cv_bridge.CvBridge()
        
        self.image_sub = self.create_subscription(CompressedImage, '/image/compressed', self.image_callback, 10)
        self.cmd_vel_pub = self.create_publisher(Twist, 'cmd_vel', 10)
        self.qr_pub = self.create_publisher(Sign, '/sign_switch', 10)
        self.pub = self.create_publisher(CompressedImage, '/camera/process_image', 10)
        self.mask_pub = self.create_publisher(CompressedImage, '/camera/mask_image', 10)
        
        self.sign_sub = self.create_subscription(Int32, 'sign_foxglove', self.sign_callback, 10)
        
        self.twist = Twist()
        self.qrDecoder = cv2.QRCodeDetector()
        
        self.is_following_line = True
        self.is_avoiding_obstacle = False
        self.external_control = False

        self.obstacle_sub = self.create_subscription(
            PerceptionTargets,
            'hobot_dnn_detection',
            self.obstacle_callback,
            10
        )
        
        self.confidence_threshold = 0.87
        self.y_axis_threshold = 300
        self.last_twist_msg = Twist()
        self.obstacle_center_x = None
        self.obstacle_direction = None

        self.avoid_angular_z = 0.0
        self.avoid_start_time = None
        self.avoid_turn_end_time = None
        self.drive_straight_end_time = None
        self.resume_direction_end_time = None
        self.final_straight_end_time = None

        # PID控制器参数
        self.kp = 0.0015  # 比例增益
        self.ki = 0.0000    # 积分增益
        self.kd = 0.001  # 微分增益

        self.prev_error = 0.0
        self.integral = 0.0

        self.create_timer(0.1, self.update)  # 添加一个定时器来定期调用update函数

    def sign_callback(self, msg):
        self.get_logger().info(f"Received sign: {msg.data}")
        if msg.data == 5:
            self.is_following_line = False
            self.external_control = True
            self.is_avoiding_obstacle = False
            self.get_logger().info("Deactivating line following due to external control.")
        elif msg.data == 6:
            self.is_following_line = True
            self.external_control = False
            self.get_logger().info("Activating line following.")

    def image_callback(self, msg):
        if not self.is_following_line or self.is_avoiding_obstacle:
            return
        
        self.get_logger().info("Image callback triggered.")
        try:
            np_arr = np.frombuffer(msg.data, np.uint8)
            image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        except cv_bridge.CvBridgeError as e:
            self.get_logger().error(f"CvBridge Error: {e}")
            return
        
        data, bbox, rectifiedImage = self.qrDecoder.detectAndDecode(image)
        if data:
            self.get_logger().info(f"QR Code detected: {data}")
            qr_msg = Sign()
            self.twist.linear.x = 0.0
            self.twist.angular.z = 0.0
            self.cmd_vel_pub.publish(self.twist)
            if data == 'ClockWise':
                qr_msg.sign_data = 3
            elif data == 'AntiClockWise':
                qr_msg.sign_data = 4
            self.get_logger().info(f'Publishing: {qr_msg.sign_data}')
            self.qr_pub.publish(qr_msg)
            return
        
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        lower_black = np.array([0, 0, 0])
        upper_black = np.array([180, 255, 50])

        mask = cv2.inRange(hsv, lower_black, upper_black)

        h, w, d = image.shape
        search_top = int(h * 7 / 11)
        search_bot = int(h / 2 + 20)
        mask[0:search_top, 0:w] = 0
        M = cv2.moments(mask)
        
        if M['m00'] > 0:
            cx = int(M['m10'] / M['m00'])
            cy = int(M['m01'] / M['m00'])
            cv2.circle(image, (cx, cy), 20, (0, 0, 255), -1)
            
            err = cx - w / 2

            # PID控制器计算
            self.integral += err
            derivative = err - self.prev_error
            self.twist.linear.x = 0.3
            self.twist.angular.z = -(self.kp * err + self.ki * self.integral + self.kd * derivative)
            self.prev_error = err

            if not self.external_control:
                self.cmd_vel_pub.publish(self.twist)

        try:
            compressed_image = CompressedImage()
            compressed_image.header.stamp = self.get_clock().now().to_msg()
            compressed_image.format = "jpeg"
            compressed_image.data = np.array(cv2.imencode('.jpg', image)[1]).tobytes()
            self.pub.publish(compressed_image)

            mask_image = CompressedImage()
            mask_image.header.stamp = self.get_clock().now().to_msg()
            mask_image.format = "jpeg"
            mask_image.data = np.array(cv2.imencode('.jpg', mask)[1]).tobytes()
            self.mask_pub.publish(mask_image)
        except cv_bridge.CvBridgeError as e:
            self.get_logger().error(f"Publish CvBridge Error: {e}")

    def obstacle_callback(self, msg):
        if self.is_avoiding_obstacle:
            return

        obstacle_detected = False
        for target in msg.targets:
            confidence = target.rois[0].confidence
            y_offset = target.rois[0].rect.y_offset
            height = target.rois[0].rect.height

            self.get_logger().info(f"检测到目标: 置信度={confidence}, y轴底部位置={y_offset + height}")
            if confidence >= self.confidence_threshold and y_offset + height >= self.y_axis_threshold:
                self.get_logger().info(f"符合条件的障碍物: 置信度={confidence}, y轴底部位置={y_offset + height}")
                if target.rois[0].type == "hindrance":
                    self.get_logger().info("调用避障逻辑")
                    self.avoid_obstacle(target)
                    self.is_avoiding_obstacle = True
                    self.avoid_start_time = time()
                    return

        if not msg.targets or not obstacle_detected:
            self.get_logger().info("未检测到目标")

    def avoid_obstacle(self, target):
        self.get_logger().info("执行避障逻辑")
        self.obstacle_center_x = target.rois[0].rect.x_offset + target.rois[0].rect.width / 2
        self.get_logger().info(f"障碍物中心点={self.obstacle_center_x}")
        if self.obstacle_center_x < 320 - 20:
            self.obstacle_direction = 'right'
            self.turn_angle('right', 0.5, 0.6)
            self.get_logger().info("向右避开障碍物")
        elif self.obstacle_center_x > 320 + 20:
            self.obstacle_direction = 'left'
            self.turn_angle('left', 0.5, 0.6)
            self.get_logger().info("向左避开障碍物")
        else:
            if self.obstacle_center_x < 320 - 10:
                self.obstacle_direction = 'right'
                self.get_logger().info("向右避开障碍物")
                self.turn_angle('right', 1.2, 0.7)
            else:
                self.obstacle_direction = 'left'
                self.get_logger().info("向左避开障碍物")
                self.turn_angle('left', 1.2, 0.7)

    def turn_angle(self, direction, ratio, turn_time):
        twist_msg = Twist()
        twist_msg.linear.x = 0.3
        self.avoid_angular_z = ratio if direction == 'left' else -ratio
        twist_msg.angular.z = self.avoid_angular_z
        self.cmd_vel_pub.publish(twist_msg)
        
        self.avoid_turn_end_time = time() + turn_time

    def drive_straight(self):
        twist_msg = Twist()
        twist_msg.linear.x = 0.3
        twist_msg.angular.z = 0.0
        self.cmd_vel_pub.publish(twist_msg)
        self.get_logger().info("正在转后直行")
        
        self.drive_straight_end_time = time() + 0.4

    def resume_direction(self):
        twist_msg = Twist()
        twist_msg.linear.x = 0.3
        resume_time = 0
        
        if self.avoid_angular_z <= 0:
            twist_msg.angular.z = -self.avoid_angular_z
            resume_time = 0.5
        elif self.avoid_angular_z > 0:
            twist_msg.angular.z = -self.avoid_angular_z
            resume_time = 0.5
    
        self.cmd_vel_pub.publish(twist_msg)
        self.get_logger().info(f"逆转恢复方向, 角速度: {twist_msg.angular.z}, 恢复时间: {resume_time}")
    
        self.resume_direction_end_time = time() + resume_time

    def final_straight(self):
        twist_msg = Twist()
        twist_msg.linear.x = 0.3
        twist_msg.angular.z = 0.0
        self.cmd_vel_pub.publish(twist_msg)
        self.get_logger().info("逆转后直线行驶")
        
        self.final_straight_end_time = time() + 0.25

    def update(self):
        current_time = time()
        if self.is_avoiding_obstacle:
            if self.avoid_turn_end_time and current_time < self.avoid_turn_end_time:
                pass
            elif self.avoid_turn_end_time and current_time >= self.avoid_turn_end_time and self.drive_straight_end_time is None:
                self.drive_straight()
            elif self.drive_straight_end_time and current_time >= self.drive_straight_end_time and self.resume_direction_end_time is None:
                self.resume_direction()
            elif self.resume_direction_end_time and current_time >= self.resume_direction_end_time and self.final_straight_end_time is None:
                self.final_straight()
            elif self.final_straight_end_time and current_time >= self.final_straight_end_time:
                self.resume_line_following()

    def resume_line_following(self):
        self.get_logger().info("恢复巡线操作")
        self.is_following_line = True
        self.is_avoiding_obstacle = False
        self.avoid_start_time = None
        self.avoid_turn_end_time = None
        self.drive_straight_end_time = None
        self.resume_direction_end_time = None
        self.final_straight_end_time = None
        self.prev_error = 0.0
        self.integral = 0.0
        self.twist.linear.x = 0.3
        self.twist.angular.z = 0.0
        self.cmd_vel_pub.publish(self.twist)
        self.get_logger().info("恢复巡线状态设置完成")

def main(args=None):
    rclpy.init(args=args)
    follower = Follower()
    rclpy.spin(follower)
    follower.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
